import math
import torch
import torch.nn as nn
import torch.nn.functional as F

# 导入概率分布类，用于强化学习中的策略采样
from torch.distributions.categorical import Categorical
from torch.distributions.multivariate_normal import MultivariateNormal
from Utils.tensor_ops import my_view


class MultiHeadAttention(nn.Module):
    """
    多头注意力机制类

    这个类实现了Transformer中使用的多头注意力机制，允许模型同时关注不同位置的信息。
    代码来源于https://github.com/wouterkool/attention-tsp/blob/master/graph_encoder.py

    多头注意力的核心思想：
    1. 将输入分成多个"头"，每个头学习不同的注意力模式
    2. 每个头独立计算注意力
    3. 将所有头的输出合并，得到最终的注意力表示

    这种机制使得模型能够同时关注序列中的不同方面，提高了模型的表达能力。
    """

    def __init__(self, n_heads, input_dim, embed_dim=None, val_dim=None, key_dim=None):
        """
        初始化多头注意力层

        Args:
            n_heads (int): 注意力头的数量
            input_dim (int): 输入特征的维度
            embed_dim (int, optional): 输出嵌入的维度，默认为None时会自动计算
            val_dim (int, optional): 值向量的维度，默认为None时会自动计算
            key_dim (int, optional): 键向量的维度，默认为None时会自动计算
        """
        super(MultiHeadAttention, self).__init__()

        # 如果没有指定值维度，则根据输出嵌入维度计算
        if val_dim is None:
            assert embed_dim is not None, "Provide either embed_dim or val_dim"
            val_dim = embed_dim // n_heads  # 每个头的值维度 = 总嵌入维度 / 头数

        # 如果没有指定键维度，则使用值维度
        if key_dim is None:
            key_dim = val_dim

        # 保存参数
        self.n_heads = n_heads           # 注意力头数量
        self.input_dim = input_dim       # 输入维度
        self.embed_dim = embed_dim       # 输出嵌入维度
        self.val_dim = val_dim           # 值向量维度
        self.key_dim = key_dim           # 键向量维度

        # 注意力缩放因子，来自论文"Attention is all you need"
        # 用于防止点积过大导致softmax梯度消失
        self.norm_factor = 1 / math.sqrt(key_dim)

        # 定义可训练的参数矩阵
        # 这些矩阵将输入转换为查询、键、值向量
        self.W_query = nn.Parameter(torch.Tensor(n_heads, input_dim, key_dim))      # 查询变换矩阵
        self.W_key = nn.Parameter(torch.Tensor(n_heads, input_dim, key_dim))        # 键变换矩阵
        self.W_val = nn.Parameter(torch.Tensor(n_heads, input_dim, val_dim))        # 值变换矩阵

        # 如果指定了输出嵌入维度，定义输出变换矩阵
        if embed_dim is not None:
            self.W_out = nn.Parameter(torch.Tensor(n_heads, key_dim, embed_dim))    # 输出变换矩阵

        # 初始化参数
        self.init_parameters()

    def init_parameters(self):
        """
        初始化网络参数

        使用均匀分布初始化所有参数，范围是[-stdv, stdv]，
        其中stdv = 1/sqrt(参数的最后一个维度大小)
        这种初始化方法有助于保持梯度的稳定性。
        """
        for param in self.parameters():
            stdv = 1. / math.sqrt(param.size(-1))
            param.data.uniform_(-stdv, stdv)

    def forward(self, q, k=None, v=None, mask=None, return_attn=False, return_attn_weight=False):
        """
        前向传播函数

        Args:
            q (torch.Tensor): 查询张量，形状为(*, n_query, input_dim)
            k (torch.Tensor, optional): 键张量，形状为(*, n_key, input_dim)，如果为None则使用q
            v (torch.Tensor, optional): 值张量，形状为(*, n_key, input_dim)，如果为None则使用k
            mask (torch.Tensor, optional): 注意力掩码，形状为(*, n_query, n_key)
            return_attn (bool): 是否返回注意力权重
            return_attn_weight (bool): 是否返回注意力权重矩阵（仅当n_heads=1时有效）

        Returns:
            torch.Tensor or tuple: 注意力输出，可能包含注意力权重
        """
        # 处理低维度输入（3维及以下）
        if q.dim() <= 3:
            out = self.forward_(q, k, v, mask, return_attn, return_attn_weight)
            if return_attn:
                out, attn = out
                assert attn.shape[0] == 1  # 确保只有一个头
                attn = attn.squeeze(0)    # 移除头的维度
                return out, attn
            return out

        # 处理高维度输入（4维及以上）
        # 将高维张量重塑为2D张量进行批处理
        hyper_dim = q.shape[:-2]  # 超越查询和键值对的维度
        q = my_view(q, [-1, *q.shape[-2:]])  # 重塑为[batch_size, n_query, input_dim]

        if k is not None:
            k = my_view(k, [-1, *k.shape[-2:]])
        if v is not None:
            v = my_view(v, [-1, *v.shape[-2:]])
        if mask is not None:
            mask = my_view(mask, [-1, *mask.shape[-2:]])

        # 执行注意力计算
        out = self.forward_(q, k, v, mask, return_attn, return_attn_weight)

        # 恢复原始维度
        if return_attn:
            out, attn = out
            if hyper_dim is not None:
                out = out.view(*hyper_dim, *out.shape[-2:])
                attn = attn.view(*hyper_dim, *attn.shape[-2:])
            return out, attn
        else:
            if hyper_dim is not None:
                out = out.view(*hyper_dim, *q.shape[-2:])
            return out

    def forward_(self, q, k=None, v=None, mask=None, return_attn=False, return_attn_weight=False):
        """
        内部前向传播函数，执行实际的注意力计算

        Args:
            q: 查询张量 (batch_size, n_query, input_dim)
            k: 键张量 (batch_size, n_key, input_dim)
            v: 值张量 (batch_size, n_key, input_dim)
            mask: 掩码张量 (batch_size, n_query, n_key)
            return_attn: 是否返回注意力权重
            return_attn_weight: 是否返回注意力权重矩阵

        Returns:
            torch.Tensor or tuple: 注意力输出，可能包含注意力权重
        """
        # 如果没有提供键和值，则使用查询进行自注意力计算
        if k is None:
            k = q  # 自注意力：键 = 查询
        if v is None:
            v = k  # 值 = 键

        # 获取张量形状信息
        batch_size, graph_size, input_dim = k.size()  # (batch_size, n_key, input_dim)
        n_query = q.size(1)                           # 查询数量

        # 形状检查
        assert q.size(0) == batch_size, "批次大小不匹配"
        assert q.size(2) == input_dim, "查询维度不匹配"
        assert input_dim == self.input_dim, "输入维度错误"

        # 将张量展平为2D，便于矩阵运算
        kflat = k.contiguous().view(-1, input_dim)   # (batch_size * n_key, input_dim)
        qflat = q.contiguous().view(-1, input_dim)   # (batch_size * n_query, input_dim)
        vflat = v.contiguous().view(-1, input_dim)   # (batch_size * n_key, input_dim)

        # 定义输出形状
        shp = (self.n_heads, batch_size, graph_size, -1)        # 键值对的形状
        shp_q = (self.n_heads, batch_size, n_query, -1)         # 查询的形状

        # 计算查询、键、值向量
        # Q = qflat * W_query -> (n_heads, batch_size, n_query, key_dim)
        Q = torch.matmul(qflat, self.W_query).view(shp_q)
        # K = kflat * W_key -> (n_heads, batch_size, n_key, key_dim)
        K = torch.matmul(kflat, self.W_key).view(shp)
        # V = vflat * W_val -> (n_heads, batch_size, n_key, val_dim)
        V = torch.matmul(vflat, self.W_val).view(shp)

        # 计算注意力分数
        # compatibility = Q * K^T * norm_factor
        # 形状: (n_heads, batch_size, n_query, n_key)
        compatibility = self.norm_factor * torch.matmul(Q, K.transpose(2, 3))

        # 如果需要返回注意力权重矩阵（仅支持单头）
        if return_attn_weight:
            assert self.n_heads == 1, "返回注意力权重只支持单头注意力"
            if mask is not None:
                mask = mask.view(1, batch_size, n_query, graph_size).expand_as(compatibility)
                compatibility[mask.bool()] = -math.inf  # 将掩码位置设为负无穷
            return compatibility.squeeze(0)  # 移除头的维度

        # 应用掩码（如果提供）
        if mask is not None:
            # 扩展掩码以匹配所有头
            mask = mask.view(1, batch_size, n_query, graph_size).expand_as(compatibility)
            compatibility[mask.bool()] = -math.inf  # 将掩码位置设为负无穷

        # 计算注意力权重
        # softmax(compatibility) -> 每个查询对每个键的注意力权重
        attn = F.softmax(compatibility, dim=-1)

        # 处理NaN值：如果有节点没有邻居，softmax会返回nan，这里将其设为0
        if mask is not None:
            attnc = attn.clone()
            attnc[mask.bool()] = 0
            attn = attnc

        # 为了解决 0*nan = nan 的问题，输入必须将V中的nan转化为0
        # 计算加权的值向量
        # heads = attn * V -> (n_heads, batch_size, n_query, val_dim)
        heads = torch.matmul(attn, V)

        # 合并多个头的输出
        # 1. 重新排列维度: (n_heads, batch_size, n_query, val_dim) -> (batch_size, n_query, n_heads, val_dim)
        # 2. 展平: -> (batch_size * n_query, n_heads * val_dim)
        # 3. 线性变换: -> (batch_size * n_query, embed_dim)
        # 4. 重塑: -> (batch_size, n_query, embed_dim)
        out = torch.mm(
            heads.permute(1, 2, 0, 3).contiguous().view(-1, self.n_heads * self.val_dim),
            self.W_out.view(-1, self.embed_dim)
        ).view(batch_size, n_query, self.embed_dim)

        # 根据需要返回注意力权重
        if return_attn:
            return out, attn
        return out


class SimpleAttention(nn.Module):
    """
    简单注意力机制类

    这是一个简化版的注意力机制，只使用单头注意力。
    相比多头注意力，计算量更小，适合简单的应用场景。
    """

    def __init__(self, h_dim):
        """
        初始化简单注意力层

        Args:
            h_dim (int): 隐藏层维度，也是输入和输出的维度
        """
        super().__init__()
        # 定义可训练的参数矩阵
        self.W_query = nn.Parameter(torch.Tensor(h_dim, h_dim))  # 查询变换矩阵
        self.W_key = nn.Parameter(torch.Tensor(h_dim, h_dim))    # 键变换矩阵
        self.W_val = nn.Parameter(torch.Tensor(h_dim, h_dim))    # 值变换矩阵
        # 初始化参数
        self.init_parameters()

    def init_parameters(self):
        """
        初始化网络参数

        使用均匀分布初始化参数，保持梯度的稳定性。
        """
        for param in self.parameters():
            stdv = 1. / math.sqrt(param.size(-1))
            param.data.uniform_(-stdv, stdv)

    def forward(self, k, q, v, mask=None):
        """
        前向传播

        Args:
            k (torch.Tensor): 键张量
            q (torch.Tensor): 查询张量
            v (torch.Tensor): 值张量
            mask (torch.Tensor, optional): 注意力掩码

        Returns:
            torch.Tensor: 注意力输出
        """
        # 计算查询、键、值向量
        Q = torch.matmul(q, self.W_query)  # 查询向量
        K = torch.matmul(k, self.W_key)    # 键向量
        V = torch.matmul(v, self.W_val)    # 值向量

        # 计算注意力分数
        # 兼容性分数 = Q * K^T / sqrt(维度)
        norm_factor = 1 / math.sqrt(Q.shape[-1])
        compat = norm_factor * torch.matmul(Q, K.transpose(-1, -2))

        # 应用掩码（如果提供）
        if mask is not None:
            compat[mask.bool()] = -math.inf  # 将掩码位置设为负无穷

        # 计算注意力权重
        # 使用torch.nan_to_num处理可能的NaN值，将NaN替换为0
        # 这是为了解决 0*nan = nan 的问题，输入必须将V中的nan转化为0
        score = torch.nan_to_num(F.softmax(compat, dim=-1), 0)

        # 计算加权的值向量
        return torch.matmul(score, V)


# 使用示例和说明：
"""
多头注意力机制的使用方法：

1. 创建多头注意力层：
   attention = MultiHeadAttention(
       n_heads=8,           # 8个注意力头
       input_dim=512,       # 输入特征维度512
       embed_dim=512        # 输出嵌入维度512
   )

2. 前向传播：
   # 自注意力：查询、键、值都来自同一个序列
   output = attention(input_sequence)

   # 交叉注意力：查询来自一个序列，键值来自另一个序列
   output = attention(
       q=query_sequence,
       k=key_sequence,
       v=value_sequence
   )

3. 带掩码的注意力：
   # 创建注意力掩码（例如用于序列填充）
   mask = torch.zeros(batch_size, seq_len, seq_len)
   mask[:, padding_len:, :] = -float('inf')

   output = attention(input_sequence, mask=mask)

4. 获取注意力权重：
   output, attn_weights = attention(input_sequence, return_attn=True)

简单注意力机制的使用方法：

1. 创建简单注意力层：
   attention = SimpleAttention(h_dim=256)

2. 前向传播：
   output = attention(k=key_sequence, q=query_sequence, v=value_sequence)

应用场景：
- 多头注意力：适用于复杂的序列建模任务，如机器翻译、文本摘要等
- 简单注意力：适用于轻量级应用，或作为复杂模型的一部分
"""